#!/usr/bin/env python

import psycopg2
import sblab
import argparse
import sys
import hashlib
import os
import gzip 
import bz2

parser = argparse.ArgumentParser(description= """
DESCRIPTION

    Upload to table sblab.main.fastqfiles the fastqfile passed as argument.
    Library_id is extracted from file name by getting the string from the begining
    of the filename till the first dot excluded.
    Quality encoding is determined by sblab.get_fastq_encoding().



EXAMPLE
    fastqfile_to_sblab.py -f db002_test.fq.gz -p 'Illumina MiSeq' -s ''

""", formatter_class= argparse.RawTextHelpFormatter)

parser.add_argument('--fastq', '-f',
                   required= True,
                   help='''Fastq file to upload.
                   ''')

parser.add_argument('--sequencer_id', '-s',
                   required= True,
                   help='''Sequencer ID for column fastqfiles.sequencer_id. Use
empty sting '' if you really don't want to assign a sequencer ID.
See `SELECT * FROM sequencers;` for current sequencers

                   ''')

parser.add_argument('--sequencing_platform', '-p',
                   required= True,
                   help='''Sequencing platform for column fastqfiles.sequencing_platform.
Use empty sting '' if you really don't want to assign a sequencer ID.
See `SELECT * FROM sequencing_platforms;`
                   ''')

parser.add_argument('--nocommit',
                   action= 'store_true',
                   help='''Do not commit changes to database, just test execution.
                   ''')
# -----------------------------------------------------------------------------
def sumfile(fobj):
    '''Returns an md5 hash for an object with read() method.'''
    m = hashlib.md5()
    while True:
        d = fobj.read(8096)
        if not d:
            break
        m.update(d)
    return m.hexdigest()


def md5sum(fname):
    '''Returns an md5 hash for file fname, or stdin if fname is "-".'''
    if fname == '-':
        ret = sumfile(sys.stdin)
    else:
        try:
            f = open(fname, 'rb')
        except:
            return 'Failed to open file %s' %fname
        ret = sumfile(f)
        f.close()
    return ret

def getFirstRead(fqfile):
    '''Get the first read, i.e. the first four lines, of fastqfile
    '''
    if fqfile.endswith('.gz'):
        fin= gzip.open(fqfile)
    elif fqfile.endswith('.bz2'):
        fin= bz2.BZ2File(fqfile)
    else:
        fin= open(fqfile)
    xread= fin.readline()
    xread += fin.readline()
    xread += fin.readline()
    xread += fin.readline().strip()
    fin.close()
    return(xread)

def fastqfilesColnames(cur):
    cur.execute('select * from fastqfiles where 1=2')
    colnames = [desc[0] for desc in cur.description]
    colnamesDict= {}
    for x in colnames:
        colnamesDict[x]= None
    return(colnamesDict)

def create_tmp_table(cursor, table):
    """Create an empty table like the template 'table'.
    cursor is an open cursor.
    """
    sql= "CREATE TEMP TABLE %s_tmp AS (SELECT * FROM %s WHERE 1=2)" %(table, table)
    print(cursor.mogrify(sql))
    cursor.execute(sql) ## Where to put results

def insert_dictionary(cur, sql, sql_dict):
    """Given open cursor 'cur' will execute 'sql' using parameters in the dict
    'sql_dict' (where key:values are column_name:value_to_insert)
    """
    print(cur.mogrify(sql, sql_dict))
    cur.execute(sql, sql_dict)    

# -----------------------------------------------------------------------------

def main():
    args = parser.parse_args()
    if args.sequencing_platform == '':
        args.sequencing_platform= None
    if args.sequencer_id == '':
        args.sequencer_id= None
    conn= sblab.get_psycopgpass()
    conn= psycopg2.connect(conn)
    cur= conn.cursor()

    columns= fastqfilesColnames(cur)
    columns['fastqfile']= os.path.split(args.fastq)[1]
    columns['encoding']= sblab.get_fastq_encoding(args.fastq)
    columns['library_id']= columns['fastqfile'].split('.')[0]
    columns['sequencing_platform']= args.sequencing_platform
    columns['sequencer_id']= args.sequencer_id
    columns['read']= getFirstRead(args.fastq)


    print('\nCalculating md5sum for "%s"...\n' %(args.fastq))
    columns['md5sum']= md5sum(args.fastq)

    create_tmp_table(cur, 'fastqfiles')    

    sql= '''\nINSERT INTO fastqfiles_tmp (fastqfile, encoding, md5sum, library_id, sequencing_platform, sequencer_id, read)
VALUES (%(fastqfile)s, %(encoding)s, %(md5sum)s, %(library_id)s, %(sequencing_platform)s, %(sequencer_id)s, %(read)s)
        '''
    insert_dictionary(cur, sql, columns)
    
    print(cur.mogrify(sql, columns))
    cur.execute(sql, columns)    
    
    sql= "INSERT INTO fastqfiles SELECT DISTINCT * FROM fastqfiles_tmp EXCEPT SELECT * FROM fastqfiles; DROP TABLE fastqfiles_tmp"
    print(sql)
    cur.execute(sql) ## Import only rows not already present 

    if not args.nocommit:
        conn.commit()
    else:
        print('\nChanges have not been committed\n')
    cur.close()
    conn.commit()
    
if __name__ == '__main__':
    main()
    sys.exit()
